from keras import Input
from keras.layers import MaxPooling2D

from .base_layer import BaseLayer
from entity.model_describe import ModelDescribe


class MaxPooling2DLayer(BaseLayer):
    def persistence(self, model_list_line: list, train_id: int) -> ModelDescribe:
        m = super().persistence(model_list_line, train_id)
        pool_size = 2 if len(model_list_line) < 2 else model_list_line[1]
        m.var1 = pool_size
        return m

    def transfer(self, model_list_line: list, inputs: Input, models: list = None):
        # 只写了MaxPooling2D没有写pool_size，用默认2
        pool_size = 2 if len(model_list_line) < 2 else model_list_line[1]
        return MaxPooling2D(pool_size=(pool_size, pool_size))(inputs)